extern crate proc_macro;
use proc_macro::TokenStream;
use quote::{format_ident, quote, ToTokens};
use syn;

/// init log, useage:
/// #[codegen::log] or #[codegen::log(info)]
///
/// default log level is info, support log level:
/// trace, debug, info, warn, error
#[proc_macro_attribute]
pub fn log(args: TokenStream, item: TokenStream) -> TokenStream {
    let mut log_level = "INFO".to_string();
    let item_cp = args.clone();
    let attributes = syn::parse_macro_input!(item_cp as syn::AttributeArgs);
    if attributes.len() > 1 {
        panic!("allowed args: error, warn, info, debug, trace");
    } else if attributes.len() == 1 {
        log_level = judge_log_level(args);
    }
    insert_log_init_code(item, log_level)
}

fn judge_log_level(args: TokenStream) -> String {
    let log_level = syn::parse::<syn::Ident>(args)
        .unwrap()
        .to_string()
        .to_uppercase();
    let support_log_levels = vec!["TRACE", "DEBUG", "INFO", "WARN", "ERROR"];
    if !support_log_levels.contains(&log_level.as_str()) {
        panic!("allowed log level: {}", support_log_levels.join(", "));
    }
    log_level
}

fn insert_log_init_code(item: TokenStream, log_level: String) -> TokenStream {
    let log_level_expr = format_ident!("{}", log_level);
    let with_log_init_code: TokenStream = (quote! {
        tracing_subscriber::fmt()
            .with_test_writer()
            .with_max_level(tracing::Level::#log_level_expr)
            .init();
    })
    .into();
    let mut asd = syn::parse::<syn::ItemFn>(item.clone()).unwrap();
    asd.block
        .stmts
        .insert(0, syn::parse(with_log_init_code).unwrap());
    return asd.into_token_stream().into();
}

#[proc_macro_attribute]
pub fn cli_log(_: TokenStream, item: TokenStream) -> TokenStream {
    let with_log_init_code: TokenStream = (quote! {
        tracing_subscriber::fmt()
            .with_max_level(tracing::Level::INFO)
            .with_test_writer()
            .with_level(false)
            .without_time()
            .with_target(false)
            .init();
    })
    .into();

    // 日志初始化操作插入到方法块的开头
    let mut asd = syn::parse::<syn::ItemFn>(item.clone()).unwrap();
    asd.block
        .stmts
        .insert(0, syn::parse(with_log_init_code).unwrap());
    asd.into_token_stream().into()
}
